package com.slwu.configure;

import com.google.common.collect.Range;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;

import java.util.*;

/**
 * 地块分表策略
 * @author wushaolin
 */
public class PlotShardingAlgorithm implements StandardShardingAlgorithm<String> {

    private static final Map<String, Integer> PROVINCE_MAPPING = new HashMap<>();

    static {
        // 省级行政区划代码映射
        PROVINCE_MAPPING.put("11", 0);  // 北京
        PROVINCE_MAPPING.put("12", 1);  // 天津
//        PROVINCE_MAPPING.put("13", 2);  // 河北
//        PROVINCE_MAPPING.put("14", 3);  // 山西
//        PROVINCE_MAPPING.put("15", 4);  // 内蒙古
//        PROVINCE_MAPPING.put("21", 5);  // 辽宁
//        PROVINCE_MAPPING.put("22", 6);  // 吉林
//        PROVINCE_MAPPING.put("23", 7);  // 黑龙江
//        PROVINCE_MAPPING.put("31", 8);  // 上海
//        PROVINCE_MAPPING.put("32", 9);  // 江苏
//        PROVINCE_MAPPING.put("33", 10); // 浙江
//        PROVINCE_MAPPING.put("34", 11); // 安徽
//        PROVINCE_MAPPING.put("35", 12); // 福建
//        PROVINCE_MAPPING.put("36", 13); // 江西
//        PROVINCE_MAPPING.put("37", 14); // 山东
//        PROVINCE_MAPPING.put("41", 15); // 河南
//        PROVINCE_MAPPING.put("42", 16); // 湖北
//        PROVINCE_MAPPING.put("43", 17); // 湖南
//        PROVINCE_MAPPING.put("44", 18); // 广东
//        PROVINCE_MAPPING.put("45", 19); // 广西
//        PROVINCE_MAPPING.put("46", 20); // 海南
//        PROVINCE_MAPPING.put("50", 21); // 重庆
//        PROVINCE_MAPPING.put("51", 22); // 四川
//        PROVINCE_MAPPING.put("52", 23); // 贵州
//        PROVINCE_MAPPING.put("53", 24); // 云南
//        PROVINCE_MAPPING.put("54", 25); // 西藏
//        PROVINCE_MAPPING.put("61", 26); // 陕西
//        PROVINCE_MAPPING.put("62", 27); // 甘肃
//        PROVINCE_MAPPING.put("63", 28); // 青海
//        PROVINCE_MAPPING.put("64", 29); // 宁夏
//        PROVINCE_MAPPING.put("65", 30); // 新疆
//        PROVINCE_MAPPING.put("71", 31); // 台湾
//        PROVINCE_MAPPING.put("81", 32); // 香港
//        PROVINCE_MAPPING.put("82", 33); // 澳门
    }

    @Override
    public String doSharding(Collection<String> availableTargetNames, PreciseShardingValue<String> preciseShardingValue) {
        String village = preciseShardingValue.getValue();
        String province = village.substring(0, 2);

        Integer tableIndex = PROVINCE_MAPPING.get(province);
        if (tableIndex == null) {
            // 默认表，处理异常数据
            tableIndex = -1;
        }

        String targetTable = tableIndex >= 0 ? preciseShardingValue.getLogicTableName() + "_" + tableIndex : preciseShardingValue.getLogicTableName();
        if (availableTargetNames.contains(targetTable)) {
            return targetTable;
        }

        throw new IllegalArgumentException("未找到对应的分表: " + targetTable);
    }

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, RangeShardingValue<String> rangeShardingValue) {
        // 处理范围查询，如查询多个省份的数据
        Range<String> range = rangeShardingValue.getValueRange();
        Set<String> result = new LinkedHashSet<>();

        if (range.hasLowerBound() && range.hasUpperBound()) {
            String lower = range.lowerEndpoint();
            String upper = range.upperEndpoint();

            // 获取范围内的所有省份
            Set<Integer> involvedTables = getInvolvedTables(lower, upper);
            for (Integer tableIndex : involvedTables) {
                String tableName = rangeShardingValue.getLogicTableName() + "_" + tableIndex;
                if (availableTargetNames.contains(tableName)) {
                    result.add(tableName);
                }
            }
        } else {
            // 如果没有明确范围，返回所有表
            return availableTargetNames;
        }

        return result;
    }

    private Set<Integer> getInvolvedTables(String lowerVillage, String upperVillage) {
        Set<Integer> tables = new HashSet<>();
        String lowerProvince = lowerVillage.substring(0, 2);
        String upperProvince = upperVillage.substring(0, 2);

        // 简化处理：返回两个省份之间的所有表
        // 实际应该根据行政区划代码的顺序计算
        boolean startCollecting = false;
        for (Map.Entry<String, Integer> entry : PROVINCE_MAPPING.entrySet()) {
            if (entry.getKey().equals(lowerProvince)) {
                startCollecting = true;
            }
            if (startCollecting) {
                tables.add(entry.getValue());
            }
            if (entry.getKey().equals(upperProvince)) {
                break;
            }
        }

        return tables;
    }

    @Override
    public Properties getProps() {
        return null;
    }

    @Override
    public void init(Properties properties) {

    }
}
